import time
import threading
import psutil
from datetime import datetime
from pynput import mouse, keyboard
from collections import defaultdict
import sys


class WorkMonitor:
    def __init__(self):
        self.data = {
            'keystrokes': 0,
            'mouse_distance': 0,
            'applications': {},
            'app_sessions': {},  # 新增：记录每个应用的使用会话
            'start_time': time.time(),
            'total_work_time': 0,
            'last_mouse_pos': (0, 0),
            'current_app': None,
            'app_start_time': time.time()
        }

        self.running = False
        self.mouse_listener = None
        self.keyboard_listener = None
        self.app_monitor_thread = None
        self.lock = threading.Lock()

        # 添加数据缓存，减少频繁计算
        self._data_cache = None
        self._cache_time = 0
        self._cache_duration = 1.0  # 缓存1秒

    def start(self):
        """启动监控"""
        self.running = True

        try:
            # 启动键盘监听
            self.keyboard_listener = keyboard.Listener(
                on_press=self.on_key_press,
                suppress=False  # 不抑制按键
            )
            self.keyboard_listener.start()

            # 启动鼠标监听
            self.mouse_listener = mouse.Listener(
                on_move=self.on_mouse_move,
                suppress=False  # 不抑制鼠标
            )
            self.mouse_listener.start()

            # 启动应用监控线程
            self.app_monitor_thread = threading.Thread(target=self.monitor_applications, daemon=True)
            self.app_monitor_thread.start()

            print("工作量监控已启动")

        except Exception as e:
            print(f"启动监控失败: {e}")
            self.running = False

    def stop(self):
        """停止监控"""
        self.running = False

        try:
            if self.keyboard_listener:
                self.keyboard_listener.stop()

            if self.mouse_listener:
                self.mouse_listener.stop()

        except Exception as e:
            print(f"停止监控时出错: {e}")

        print("工作量监控已停止")

    def on_key_press(self, key):
        """键盘按键回调"""
        try:
            with self.lock:
                self.data['keystrokes'] += 1
                self._invalidate_cache()
        except Exception as e:
            pass  # 忽略键盘监听错误

    def on_mouse_move(self, x, y):
        """鼠标移动回调"""
        try:
            with self.lock:
                last_x, last_y = self.data['last_mouse_pos']
                if last_x != 0 or last_y != 0:
                    distance = ((x - last_x) ** 2 + (y - last_y) ** 2) ** 0.5
                    self.data['mouse_distance'] += distance
                self.data['last_mouse_pos'] = (x, y)
                self._invalidate_cache()
        except Exception as e:
            pass  # 忽略鼠标监听错误

    def monitor_applications(self):
        """监控活动应用程序"""
        while self.running:
            try:
                current_app = self.get_active_window()
                current_time = time.time()

                with self.lock:
                    # 如果应用程序发生变化
                    if self.data['current_app'] and self.data['current_app'] != current_app:
                        # 记录上一个应用的结束时间
                        app_name = self.data['current_app']
                        session_duration = current_time - self.data['app_start_time']

                        # 更新应用总时间
                        if app_name not in self.data['applications']:
                            self.data['applications'][app_name] = {
                                'name': app_name,
                                'total_time': 0,
                                'sessions': 0,
                                'last_start_time': self.data['app_start_time'],
                                'last_end_time': current_time,
                                'current_session_start': None
                            }

                        self.data['applications'][app_name]['total_time'] += session_duration
                        self.data['applications'][app_name]['sessions'] += 1
                        self.data['applications'][app_name]['last_end_time'] = current_time
                        self.data['applications'][app_name]['current_session_start'] = None

                        # 记录会话历史
                        if app_name not in self.data['app_sessions']:
                            self.data['app_sessions'][app_name] = []

                        self.data['app_sessions'][app_name].append({
                            'start_time': self.data['app_start_time'],
                            'end_time': current_time,
                            'duration': session_duration
                        })

                    # 更新当前应用
                    if current_app != self.data['current_app']:
                        self.data['current_app'] = current_app
                        self.data['app_start_time'] = current_time

                        # 如果是新应用，初始化记录
                        if current_app and current_app not in self.data['applications']:
                            self.data['applications'][current_app] = {
                                'name': current_app,
                                'total_time': 0,
                                'sessions': 0,
                                'last_start_time': current_time,
                                'last_end_time': None,
                                'current_session_start': current_time
                            }
                        elif current_app:
                            # 更新开始时间（如果是重新使用该应用）
                            self.data['applications'][current_app]['last_start_time'] = current_time
                            self.data['applications'][current_app]['last_end_time'] = None
                            self.data['applications'][current_app]['current_session_start'] = current_time

                        self._invalidate_cache()

                    # 更新总工作时间
                    self.data['total_work_time'] = current_time - self.data['start_time']

                time.sleep(2)  # 每2秒检查一次，减少CPU占用

            except Exception as e:
                print(f"应用监控错误: {e}")
                time.sleep(2)

    def get_active_window(self):
        """获取当前活动窗口的应用名称"""
        try:
            if sys.platform == "win32":
                try:
                    import win32gui
                    import win32process

                    hwnd = win32gui.GetForegroundWindow()
                    if hwnd:
                        _, pid = win32process.GetWindowThreadProcessId(hwnd)
                        process = psutil.Process(pid)
                        return process.name()
                except ImportError:
                    pass

            elif sys.platform == "darwin":  # macOS
                try:
                    from AppKit import NSWorkspace
                    active_app = NSWorkspace.sharedWorkspace().activeApplication()
                    return active_app['NSApplicationName']
                except ImportError:
                    pass

            elif sys.platform.startswith("linux"):
                try:
                    import subprocess
                    result = subprocess.run(['xdotool', 'getwindowfocus', 'getwindowname'],
                                            capture_output=True, text=True, timeout=1)
                    if result.returncode == 0:
                        return result.stdout.strip()
                except (subprocess.SubprocessError, FileNotFoundError, subprocess.TimeoutExpired):
                    pass

            # 备用方案：返回最高CPU使用率的进程
            processes = list(psutil.process_iter(['pid', 'name', 'cpu_percent']))
            if processes:
                top_process = max(processes, key=lambda p: p.info['cpu_percent'] or 0)
                return top_process.info['name']

        except Exception as e:
            pass

        return "Unknown"

    def _invalidate_cache(self):
        """使缓存失效"""
        self._data_cache = None

    def get_data(self):
        """获取当前监控数据（带缓存）"""
        current_time = time.time()

        # 检查缓存是否有效
        if (self._data_cache is not None and
                current_time - self._cache_time < self._cache_duration):
            return self._data_cache

        with self.lock:
            # 更新当前应用的时间
            data_copy = self.data.copy()
            data_copy['applications'] = {k: v.copy() for k, v in self.data['applications'].items()}
            data_copy['app_sessions'] = {k: v.copy() for k, v in self.data['app_sessions'].items()}

            # 如果有当前应用，计算当前会话时间
            if self.data['current_app']:
                app_name = self.data['current_app']
                if app_name not in data_copy['applications']:
                    data_copy['applications'][app_name] = {
                        'name': app_name,
                        'total_time': 0,
                        'sessions': 0,
                        'last_start_time': self.data['app_start_time'],
                        'last_end_time': None,
                        'current_session_start': self.data['app_start_time']
                    }

                current_session_time = current_time - self.data['app_start_time']
                data_copy['applications'][app_name]['total_time'] += current_session_time
                # 正在使用的应用，结束时间为None
                data_copy['applications'][app_name]['last_end_time'] = None
                data_copy['applications'][app_name]['current_session_start'] = self.data['app_start_time']

            # 更新总工作时间
            data_copy['total_work_time'] = current_time - self.data['start_time']

            # 更新缓存
            self._data_cache = data_copy
            self._cache_time = current_time

            return data_copy

    def reset_data(self):
        """重置所有数据"""
        with self.lock:
            self.data = {
                'keystrokes': 0,
                'mouse_distance': 0,
                'applications': {},
                'app_sessions': {},
                'start_time': time.time(),
                'total_work_time': 0,
                'last_mouse_pos': (0, 0),
                'current_app': None,
                'app_start_time': time.time()
            }
            self._invalidate_cache()
        print("数据已重置")


# 测试代码
if __name__ == "__main__":
    monitor = WorkMonitor()
    monitor.start()

    try:
        while True:
            time.sleep(5)
            data = monitor.get_data()
            print(f"按键: {data['keystrokes']}, 鼠标距离: {data['mouse_distance']:.1f}px")
            print(f"当前应用: {data['current_app']}")
            print(f"工作时间: {data['total_work_time']:.1f}s")
            print("---")
    except KeyboardInterrupt:
        monitor.stop()
